#include "DisassemblyPopup.h"

#include <QPainter>

#include "CachedFontMetrics.h"
#include "Configuration.h"
#include "MiscUtil.h"
#include "StringUtil.h"

DisassemblyPopup::DisassemblyPopup(QWidget* parent)
    : QFrame(parent, Qt::Tool | Qt::FramelessWindowHint |
                         Qt::WindowStaysOnTopHint |
                         Qt::WindowDoesNotAcceptFocus),
      mFontMetrics(nullptr),
      mDisasm(ConfigUint("Disassembler", "MaxModuleSize")) {
  addr = 0;
  addrText = nullptr;
  connect(Config(), SIGNAL(fontsUpdated()), this, SLOT(updateFont()));
  connect(Config(), SIGNAL(colorsUpdated()), this, SLOT(updateColors()));
  connect(Config(), SIGNAL(tokenizerConfigUpdated()), this,
          SLOT(tokenizerConfigUpdated()));
  updateFont();
  updateColors();
  setFrameStyle(QFrame::Panel);
  setLineWidth(2);
  mMaxInstructions = 20;
}

void DisassemblyPopup::updateColors() {
  disassemblyBackgroundColor = ConfigColor("DisassemblyBackgroundColor");
  disassemblyTracedColor = ConfigColor("DisassemblyTracedBackgroundColor");
  labelColor = ConfigColor("DisassemblyLabelColor");
  labelBackgroundColor = ConfigColor("DisassemblyLabelBackgroundColor");
  commentColor = ConfigColor("DisassemblyCommentColor");
  commentBackgroundColor = ConfigColor("DisassemblyCommentBackgroundColor");
  commentAutoColor = ConfigColor("DisassemblyAutoCommentColor");
  commentAutoBackgroundColor =
      ConfigColor("DisassemblyAutoCommentBackgroundColor");
  QPalette palette;
  palette.setColor(QPalette::Foreground,
                   ConfigColor("AbstractTableViewSeparatorColor"));
  setPalette(palette);
}

void DisassemblyPopup::tokenizerConfigUpdated() { mDisasm.UpdateConfig(); }

void DisassemblyPopup::updateFont() {
  delete mFontMetrics;
  setFont(ConfigFont("Disassembly"));
  QFontMetricsF metrics(font());
  mFontMetrics = new CachedFontMetrics(this, font());
  // Update font size, used in layout calculations.
  charWidth = mFontMetrics->width('W');
  charHeight = metrics.height();
}

void DisassemblyPopup::paintEvent(QPaintEvent* event) {
  QRect viewportRect(0, 0, width(), height());
  QPainter p(this);
  p.setFont(font());

  // Render background
  p.fillRect(viewportRect, disassemblyBackgroundColor);
  // Draw Address
  p.setPen(QPen(labelColor));
  int addrWidth = mFontMetrics->width(addrText);
  p.fillRect(3, 2 + lineWidth(), addrWidth, charHeight,
             QBrush(labelBackgroundColor));
  p.drawText(3, 2, addrWidth, charHeight, 0, addrText);
  // Draw Comments
  if (!addrComment.isEmpty()) {
    int commentWidth = mFontMetrics->width(addrComment);
    QBrush background = QBrush(addrCommentAuto ? commentAutoBackgroundColor
                                               : commentBackgroundColor);
    p.setPen(addrCommentAuto ? commentAutoColor : commentColor);
    p.fillRect(3 + addrWidth, 2, commentWidth, charHeight, background);
    p.drawText(3 + addrWidth, 2, commentWidth, charHeight, 0, addrComment);
  }
  // Draw Instructions
  int y = charHeight + 1;
  for (auto& instruction : mDisassemblyToken) {
    if (instruction.second)
      p.fillRect(QRect(3, y, mWidth - 3, charHeight), disassemblyTracedColor);
    RichTextPainter::paintRichText(&p, 3, y, mWidth - 3, charHeight, 0,
                                   instruction.first, mFontMetrics);
    y += charHeight;
  }
  QFrame::paintEvent(event);
}

void DisassemblyPopup::setAddress(duint Address) {
  addr = Address;
  QList<Instruction_t> instBuffer;
  mDisassemblyToken.clear();
  if (addr != 0) {
    mWidth = 1;
    // Get RVA
    auto addr = Address;
    duint size;
    duint base = DbgMemFindBaseAddr(addr, &size);
    // Prepare RVA of every instruction
    unsigned int i = 0;
    instBuffer.clear();
    auto nextAddr = addr;
    bool hadBranch = false;
    duint bestBranch = 0;
    byte data[64];
    do {
      if (nextAddr >= base + size) break;
      if (!DbgMemRead(nextAddr, data, sizeof(data))) break;
      auto instruction = mDisasm.DisassembleAt(data, sizeof(data), 0, nextAddr);
      if (!instruction.length) break;
      instBuffer.append(std::move(instruction));
      if (!hadBranch || bestBranch <= nextAddr) {
        if (instruction.instStr.contains("ret")) break;
        if (instruction.instStr.contains("jmp") &&
            instruction.instStr.contains("["))
          break;
      }
      if (instruction.branchDestination &&
          !instruction.instStr.contains("call") &&
          !instruction.instStr.contains("ret")) {
        hadBranch = true;
        if (instruction.branchDestination > bestBranch)
          bestBranch = instruction.branchDestination;
      }
      auto nextAddr2 = nextAddr + instruction.length;
      if (nextAddr2 == nextAddr)
        break;
      else
        nextAddr = nextAddr2;
      if (DbgGetFunctionTypeAt(nextAddr - 1) == FUNC_END) break;
      i++;
    } while (i < mMaxInstructions);
    // Disassemble
    for (auto& instruction : instBuffer) {
      RichTextPainter::List richText;
      ZydisTokenizer::TokenToRichText(instruction.tokens, richText, nullptr);
      // Calculate width
      int currentInstructionWidth = 0;
      for (auto& token : richText)
        currentInstructionWidth += mFontMetrics->width(token.text);
      mWidth = std::max(mWidth, currentInstructionWidth);
      mDisassemblyToken.push_back(std::make_pair(
          std::move(richText),
          DbgFunctions()->GetTraceRecordHitCount(instruction.rva) != 0));
    }
    // Address
    addrText = getSymbolicName(addr);
    // Comments
    GetCommentFormat(addr, addrComment, &addrCommentAuto);
    if (addrComment.length()) addrText.append(' ');
    // Calculate width of address
    mWidth = std::max(mWidth, mFontMetrics->width(addrText) +
                                  mFontMetrics->width(addrComment));
    mWidth += charWidth * 6;
    // Resize popup
    resize(mWidth + 2, charHeight * int(mDisassemblyToken.size() + 1) + 2);
  }
  update();
}

duint DisassemblyPopup::getAddress() { return addr; }

void DisassemblyPopup::hide() {
  addr = 0;
  QFrame::hide();
}
